#!/usr/bin/env python

'''
Sends sniffed IEEE 802.15.4 packets to Wireshark via a named pipe.
(ryan@riverloopsecurity.com)
'''

import sys
import os
import argparse
import subprocess

from killerbee import *

class ShowDevAction(argparse.Action):
    def __call__(self, *a, **kw):
        show_dev()
        sys.exit(0)

def parse_args():
    global args

    # Command-line arguments
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('-i', '--iface', '--dev', action='store', dest='devstring',
            help='Device to use for sniffing')
    parser.add_argument('-p', '--ppi', action='store_true',
            help='Include CACE Per-Packet Information in PCAP')
    parser.add_argument('-c', '-f', '--channel', action='store', type=int, required=True,
            help='Channel on which to sniff')
    parser.add_argument('-s', '--subghz_page', action='store', type=int, required=False, default=0,
            help='SubGHz page on which to sniff')
    parser.add_argument('-n', '--count', action='store', type=int, default=None,
            metavar='COUNT', help='Limit capture to COUNT packets')
    parser.add_argument('-D', action = ShowDevAction, nargs=0,
            help='Show device info and exit')
    args = parser.parse_args()

def start_wireshark():
    spargs = dict(
        args = ['wireshark','-k','-i','-'],    # Read packets from stdin immediately
        stdin = subprocess.PIPE,
        stderr = open(os.devnull, 'w'),
    )

    # Put Wireshark in its own process group.
    # On *nix, this is necessary to keep Wireshark open after this script exists.
    # It doesn't appear to be necessary on Windows to accomplish the same,
    # but seems like a good idea to take the equivalent action on both platforms..
    if os.name == 'posix':
        spargs['preexec_fn'] = os.setpgrp
    elif os.name == 'nt':
        spargs['creationflags'] = subprocess.CREATE_NEW_PROCESS_GROUP

    return subprocess.Popen(**spargs)

def main():
    parse_args()

    # Start KillerBee
    with KillerBee(device=args.devstring) as kb:
        try:
            kb.set_channel(args.channel, args.subghz_page)
        except ValueError as e:
            print 'ERROR:', e
            sys.exit(1)

        kb.sniffer_on()

        # Start Wireshark
        wireshark_proc = start_wireshark()

        # Create a PCAP dumper to write packets to wireshark
        with PcapDumper(DLT_IEEE802_15_4, wireshark_proc.stdin, ppi=args.ppi) as pd:

            #rf_freq_mhz = (args.channel - 10) * 5 + 2400
            #print("zbwireshark: listening on \'{0}\'".format(kb.get_dev_info()[0]))
            rf_freq_mhz = kb.frequency(args.channel, args.subghz_page) / 1000.0
            print("zbwireshark: listening on \'{0}\', channel {1}, page {2} ({3} MHz), link-type DLT_IEEE802_15_4, capture size 127 bytes".format(kb.get_dev_info()[0], args.channel, args.subghz_page, rf_freq_mhz))
            try:
                packetcount = 0
                while args.count != packetcount:
                    # Wait for the next packet
                    packet = kb.pnext()

                    rc = wireshark_proc.poll()
                    if rc is not None:
                        print("Wireshark exited ({0})".format(rc))
                        break

                    if packet != None:
                        packetcount+=1
                        pd.pcap_dump(packet['bytes'], ant_dbm=packet['dbm'], freq_mhz=rf_freq_mhz)

            except KeyboardInterrupt:
                pass
            except IOError as e:
                if e.errno == 32:
                    #print("ERROR: Pipe broken. Was Wireshark closed or stopped?")
                    pass
                else:
                    raise

            kb.sniffer_off()
            print("{0} packets captured".format(packetcount))

if __name__ == '__main__':
    main()
